import numpy as np
import time, pickle
import cv2

WIDTH = 640
HEIGHT = 480

INTERNAL_WIDTH = 1280
INTERNAL_HEIGHT = 720

def camera_calibration(img, mtx, dist):
    h, w = img.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
    dst = cv2.undistort(img,mtx,dist,None,newcameramtx)
    x,y,w,h = roi
    dst1 = dst[y:y+h,x:x+w]
    return dst1

def load_camera_params():
    with open("../configs/cam_calib_mtx.bin", 'rb') as cam_calib_file:
        calib_mtx = pickle.load(cam_calib_file)
    with open("../configs/cam_calib_dist.bin", 'rb') as cam_calib_file:
        calib_dist = pickle.load(cam_calib_file)
    with open("../configs/cam_idx.bin", 'rb') as cam_idx_file:
        cam_idx_list = pickle.load(cam_idx_file)

    return cam_idx_list, calib_mtx, calib_dist

def get_capture(camera_id, cp = None):
    if cp is None:
        cp = cv2.VideoCapture(camera_id)
    else:
        cp.open(camera_id)
    cp.set(cv2.CAP_PROP_FRAME_WIDTH, INTERNAL_WIDTH)
    cp.set(cv2.CAP_PROP_FRAME_HEIGHT, INTERNAL_HEIGHT)
    return cp

def get_frame(cp):
    _, frame = cp.read()
    if not _:
        return None
    else:
        return cv2.resize(frame, (WIDTH, HEIGHT), cv2.INTER_AREA)

def get_capture():
    cp = cv2.VideoCapture()
    cp.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))
    cp.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
    cp.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)
    cp.set(cv2.CAP_PROP_FPS, 5)
    cp.set(cv2.CAP_PROP_FPS, 5)
    return cp